#pragma once

#include <pthread.h>
#include <stdlib.h>
#include <cstdio>
#include <unistd.h>

// using Function = void* (*)(void*);

struct Task
{
  void* (*func)(void*);
  void *arg;
  Task(void* (*func)(void*), void *arg)
  {
    this->func = func;
    this->arg = arg;
  }

  Task()
  {
    this->func = NULL;
    this->arg = NULL;
  }
};

class ThreadPool
{
private:
  Task *tasks;
  pthread_t *workers;
  pthread_mutex_t mutex;
  pthread_cond_t cond;

  size_t capacity;
  size_t worker_count;

  size_t front, rear;

  bool is_shutdown;

public:
  ThreadPool(size_t capacity, size_t worker_count);
  ~ThreadPool();

  static void *execute(void *arg);
  bool submit(Task task);
  void shutdown();

};

ThreadPool::ThreadPool(size_t capacity, size_t worker_count):
  capacity(capacity), worker_count(worker_count), is_shutdown(false)
{
  int err;
  this->tasks = (Task*)malloc(sizeof(Task) * capacity);
  
  this->workers = (pthread_t*)malloc(sizeof(pthread_t) * worker_count);
  
  err = pthread_mutex_init (&this->mutex, NULL);
  if (err != 0) perror ("");

  err = pthread_cond_init (&this->cond, NULL);
  if (err != 0) perror ("");
  
  this->front = this->rear = 0;

  for (size_t i = 0; i < worker_count; i++)
    pthread_create (&workers[i], NULL, ThreadPool::execute, this);
}

ThreadPool::~ThreadPool()
{
  if (this->is_shutdown == false)
    this->shutdown ();
  free (this->tasks);
  free (this->workers);
  pthread_mutex_destroy (&this->mutex);
  pthread_cond_destroy (&this->cond);
}

void *
ThreadPool::execute(void *arg)
{
  ThreadPool *pool = (ThreadPool *)arg;
  Task task;
  while (true)
  {
    pthread_mutex_lock (&pool->mutex);

    while (pool->front == pool->rear && pool->is_shutdown == false)
      pthread_cond_wait (&pool->cond, &pool->mutex);

    if (pool->front == pool->rear && pool->is_shutdown)
    {
      pthread_mutex_unlock (&pool->mutex);
      break;
    }

    task = pool->tasks[pool->front];
    pool->front = (pool->front + 1) % pool->capacity;

    pthread_mutex_unlock (&pool->mutex);

    task.func (task.arg);
  }
  pthread_exit (NULL);
}

bool
ThreadPool::submit(Task task)
{
  bool success = false;
  pthread_mutex_lock (&this->mutex);

  if (this->front != (this->rear + 1) % this->capacity
    && this->is_shutdown == false)
  {
    this->tasks[this->rear] = task;
    this->rear = (this->rear + 1) % this->capacity;
    success = true;
  }

  pthread_mutex_unlock (&this->mutex);

  pthread_cond_signal (&this->cond);

  return success;
}

void
ThreadPool::shutdown()
{
  pthread_mutex_lock (&this->mutex);

  this->is_shutdown = true;

  pthread_cond_broadcast (&this->cond);

  pthread_mutex_unlock (&this->mutex);

  for (size_t i = 0; i < this->worker_count; i++)
    pthread_join (workers[i], NULL);
}